package main

import (
    "flag"
    "log"
    "net/url"
    "os"
    "os/signal"
    "time"
    "github.com/gorilla/websocket"
)

var addr = flag.String("addr", "localhost:8080", "http service address")

type Response struct {
    Type      string      `json:"type"`
    Code      int         `json:"code"`
    Data      interface{} `json:"data"`
    MessageId string      `json:"messageId"`
}

func main() {
    flag.Parse()
    log.SetFlags(0)

    // 捕获中断信号
    interrupt := make(chan os.Signal, 1)
    signal.Notify(interrupt, os.Interrupt)

    // 构建WebSocket URL
    u := url.URL{Scheme: "ws", Host: *addr, Path: "/worktool/websocket/test-device"}
    log.Printf("connecting to %s", u.String())

    // 连接WebSocket服务器
    c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
    if err != nil {
        log.Fatal("dial:", err)
    }
    defer c.Close()

    done := make(chan struct{})

    // 启动一个goroutine来读取消息
    go func() {
        defer close(done)
        for {
            var response Response
            err := c.ReadJSON(&response)
            if err != nil {
                log.Println("read:", err)
                return
            }
            log.Printf("recv: %+v", response)
        }
    }()

    // 模拟心跳
    ticker := time.NewTicker(time.Second * 30)
    defer ticker.Stop()

    for {
        select {
        case <-done:
            return
        case <-ticker.C:
            // 发送心跳消息
            err := c.WriteJSON(map[string]interface{}{
                "type": "heartbeat",
                "data": map[string]interface{}{
                    "timestamp": time.Now().Unix(),
                },
                "messageId": time.Now().Format("20060102150405"),
            })
            if err != nil {
                log.Println("write:", err)
                return
            }
        case <-interrupt:
            log.Println("interrupt")

            // 优雅关闭连接
            err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
            if err != nil {
                log.Println("write close:", err)
                return
            }
            select {
            case <-done:
            case <-time.After(time.Second):
            }
            return
        }
    }
} 